from beancount.core import data, account, getters
from beancount.core.amount import Amount
from beancount.core.number import Decimal
from beancount.parser import parser
from beancount.loader import load_file
from datetime import datetime
import os
import json
from typing import Dict, List, Tuple, Optional, Set

# 设置默认的beancount文件路径
DEFAULT_BEANCOUNT_FILE = os.path.join(os.path.dirname(__file__), "ledger.beancount")

def get_balance_sheet(file_path: str = None, end_date: str = None) -> Dict:
    """
    生成资产负债表
    
    Args:
        file_path: beancount文件路径
        end_date: 结算日期，格式为YYYY-MM-DD
        
    Returns:
        资产负债表数据字典
    """
    if file_path is None:
        file_path = DEFAULT_BEANCOUNT_FILE
        
    try:
        # 加载beancount文件
        entries, errors, options_map = load_file(file_path)
        if errors:
            return {"success": False, "error": f"账套文件存在错误: {errors}"}
            
        # 处理日期筛选
        if end_date:
            end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").date()
            filtered_entries = [entry for entry in entries if 
                              (hasattr(entry, 'date') and entry.date <= end_date_obj)]
        else:
            filtered_entries = entries
            
        # 分类账户余额
        assets = {}
        liabilities = {}
        equity = {}
        
        # 计算每个账户的余额
        for entry in filtered_entries:
            if isinstance(entry, data.Transaction):
                for posting in entry.postings:
                    account_name = posting.account
                    account_type = account_name.split(':', 1)[0]
                    
                    # 获取金额
                    amount_decimal = posting.units.number
                    currency = posting.units.currency
                    
                    # 根据账户类型分类
                    if account_type == 'Assets':
                        if account_name not in assets:
                            assets[account_name] = {}
                        if currency not in assets[account_name]:
                            assets[account_name][currency] = Decimal('0')
                        assets[account_name][currency] += amount_decimal
                    elif account_type == 'Liabilities':
                        if account_name not in liabilities:
                            liabilities[account_name] = {}
                        if currency not in liabilities[account_name]:
                            liabilities[account_name][currency] = Decimal('0')
                        liabilities[account_name][currency] += amount_decimal
                    elif account_type == 'Equity':
                        if account_name not in equity:
                            equity[account_name] = {}
                        if currency not in equity[account_name]:
                            equity[account_name][currency] = Decimal('0')
                        equity[account_name][currency] += amount_decimal
        
        # 计算总资产和总负债
        total_assets = {}
        for acc, currencies in assets.items():
            for currency, amount in currencies.items():
                if currency not in total_assets:
                    total_assets[currency] = Decimal('0')
                total_assets[currency] += amount
                
        total_liabilities = {}
        for acc, currencies in liabilities.items():
            for currency, amount in currencies.items():
                if currency not in total_liabilities:
                    total_liabilities[currency] = Decimal('0')
                total_liabilities[currency] += amount
                
        total_equity = {}
        for acc, currencies in equity.items():
            for currency, amount in currencies.items():
                if currency not in total_equity:
                    total_equity[currency] = Decimal('0')
                total_equity[currency] += amount
        
        # 格式化结果
        formatted_assets = format_account_balances(assets)
        formatted_liabilities = format_account_balances(liabilities)
        formatted_equity = format_account_balances(equity)
        formatted_total_assets = format_currency_amounts(total_assets)
        formatted_total_liabilities = format_currency_amounts(total_liabilities)
        formatted_total_equity = format_currency_amounts(total_equity)
        
        return {
            "success": True,
            "date": end_date or "latest",
            "assets": formatted_assets,
            "liabilities": formatted_liabilities,
            "equity": formatted_equity,
            "total_assets": formatted_total_assets,
            "total_liabilities": formatted_total_liabilities,
            "total_equity": formatted_total_equity
        }
    except Exception as e:
        return {"success": False, "error": str(e)}

def get_income_statement(file_path: str = None, start_date: str = None, end_date: str = None) -> Dict:
    """
    生成利润表
    
    Args:
        file_path: beancount文件路径
        start_date: 开始日期，格式为YYYY-MM-DD
        end_date: 结束日期，格式为YYYY-MM-DD
        
    Returns:
        利润表数据字典
    """
    if file_path is None:
        file_path = DEFAULT_BEANCOUNT_FILE
        
    try:
        # 加载beancount文件
        entries, errors, options_map = load_file(file_path)
        if errors:
            return {"success": False, "error": f"账套文件存在错误: {errors}"}
            
        # 处理日期筛选
        start_date_obj = None
        end_date_obj = None
        
        if start_date:
            start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").date()
        if end_date:
            end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").date()
            
        filtered_entries = entries
        if start_date_obj or end_date_obj:
            filtered_entries = [entry for entry in entries if 
                              (hasattr(entry, 'date') and 
                               (not start_date_obj or entry.date >= start_date_obj) and
                               (not end_date_obj or entry.date <= end_date_obj))]
            
        # 分类收入和支出
        income = {}
        expenses = {}
        
        # 计算每个账户的金额
        for entry in filtered_entries:
            if isinstance(entry, data.Transaction):
                for posting in entry.postings:
                    account_name = posting.account
                    account_type = account_name.split(':', 1)[0]
                    
                    # 获取金额
                    amount_decimal = posting.units.number
                    currency = posting.units.currency
                    
                    # 根据账户类型分类
                    if account_type == 'Income':
                        if account_name not in income:
                            income[account_name] = {}
                        if currency not in income[account_name]:
                            income[account_name][currency] = Decimal('0')
                        income[account_name][currency] += amount_decimal
                    elif account_type == 'Expenses':
                        if account_name not in expenses:
                            expenses[account_name] = {}
                        if currency not in expenses[account_name]:
                            expenses[account_name][currency] = Decimal('0')
                        expenses[account_name][currency] += amount_decimal
        
        # 计算总收入和总支出
        total_income = {}
        for acc, currencies in income.items():
            for currency, amount in currencies.items():
                if currency not in total_income:
                    total_income[currency] = Decimal('0')
                total_income[currency] += amount
                
        total_expenses = {}
        for acc, currencies in expenses.items():
            for currency, amount in currencies.items():
                if currency not in total_expenses:
                    total_expenses[currency] = Decimal('0')
                total_expenses[currency] += amount
        
        # 计算净利润
        net_profit = {}
        all_currencies = set(list(total_income.keys()) + list(total_expenses.keys()))
        for currency in all_currencies:
            income_amount = total_income.get(currency, Decimal('0'))
            expense_amount = total_expenses.get(currency, Decimal('0'))
            net_profit[currency] = income_amount - expense_amount
        
        # 格式化结果
        formatted_income = format_account_balances(income)
        formatted_expenses = format_account_balances(expenses)
        formatted_total_income = format_currency_amounts(total_income)
        formatted_total_expenses = format_currency_amounts(total_expenses)
        formatted_net_profit = format_currency_amounts(net_profit)
        
        return {
            "success": True,
            "period": {
                "start_date": start_date or "earliest",
                "end_date": end_date or "latest"
            },
            "income": formatted_income,
            "expenses": formatted_expenses,
            "total_income": formatted_total_income,
            "total_expenses": formatted_total_expenses,
            "net_profit": formatted_net_profit
        }
    except Exception as e:
        return {"success": False, "error": str(e)}

def get_cash_flow_statement(file_path: str = None, start_date: str = None, end_date: str = None) -> Dict:
    """
    生成现金流量表
    
    Args:
        file_path: beancount文件路径
        start_date: 开始日期，格式为YYYY-MM-DD
        end_date: 结束日期，格式为YYYY-MM-DD
        
    Returns:
        现金流量表数据字典
    """
    if file_path is None:
        file_path = DEFAULT_BEANCOUNT_FILE
        
    try:
        # 加载beancount文件
        entries, errors, options_map = load_file(file_path)
        if errors:
            return {"success": False, "error": f"账套文件存在错误: {errors}"}
            
        # 处理日期筛选
        start_date_obj = None
        end_date_obj = None
        
        if start_date:
            start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").date()
        if end_date:
            end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").date()
            
        filtered_entries = entries
        if start_date_obj or end_date_obj:
            filtered_entries = [entry for entry in entries if 
                              (hasattr(entry, 'date') and 
                               (not start_date_obj or entry.date >= start_date_obj) and
                               (not end_date_obj or entry.date <= end_date_obj))]
        
        # 现金账户列表 (通常是Assets:Cash和Assets:Bank等)
        cash_accounts = [acc for acc in getters.get_accounts(entries) 
                        if acc.startswith('Assets:Cash') or acc.startswith('Assets:Bank')]
        
        # 分类现金流
        operating_flows = {}  # 经营活动现金流
        investing_flows = {}  # 投资活动现金流
        financing_flows = {}  # 筹资活动现金流
        
        # 分析每笔交易
        for entry in filtered_entries:
            if isinstance(entry, data.Transaction):
                # 检查是否涉及现金账户
                cash_postings = [p for p in entry.postings if p.account in cash_accounts]
                if not cash_postings:
                    continue
                    
                for cash_posting in cash_postings:
                    # 获取现金流动金额
                    amount_decimal = cash_posting.units.number
                    currency = cash_posting.units.currency
                    
                    # 确定现金流类型
                    flow_type = determine_cash_flow_type(entry, cash_posting)
                    
                    # 根据流类型分类
                    if flow_type == 'operating':
                        if currency not in operating_flows:
                            operating_flows[currency] = Decimal('0')
                        operating_flows[currency] += amount_decimal
                    elif flow_type == 'investing':
                        if currency not in investing_flows:
                            investing_flows[currency] = Decimal('0')
                        investing_flows[currency] += amount_decimal
                    elif flow_type == 'financing':
                        if currency not in financing_flows:
                            financing_flows[currency] = Decimal('0')
                        financing_flows[currency] += amount_decimal
        
        # 计算总现金流
        total_cash_flow = {}
        all_currencies = set(list(operating_flows.keys()) + 
                          list(investing_flows.keys()) + 
                          list(financing_flows.keys()))
                          
        for currency in all_currencies:
            operating_amount = operating_flows.get(currency, Decimal('0'))
            investing_amount = investing_flows.get(currency, Decimal('0'))
            financing_amount = financing_flows.get(currency, Decimal('0'))
            total_cash_flow[currency] = operating_amount + investing_amount + financing_amount
        
        # 格式化结果
        formatted_operating = format_currency_amounts(operating_flows)
        formatted_investing = format_currency_amounts(investing_flows)
        formatted_financing = format_currency_amounts(financing_flows)
        formatted_total = format_currency_amounts(total_cash_flow)
        
        return {
            "success": True,
            "period": {
                "start_date": start_date or "earliest",
                "end_date": end_date or "latest"
            },
            "operating_activities": formatted_operating,
            "investing_activities": formatted_investing,
            "financing_activities": formatted_financing,
            "total_cash_flow": formatted_total,
            "cash_accounts": cash_accounts
        }
    except Exception as e:
        return {"success": False, "error": str(e)}

def determine_cash_flow_type(transaction, cash_posting) -> str:
    """
    确定现金流类型
    
    Args:
        transaction: 交易记录
        cash_posting: 现金账户的分录
        
    Returns:
        现金流类型: 'operating', 'investing', 或 'financing'
    """
    # 默认为经营活动
    flow_type = 'operating'
    
    # 查找交易中的其他账户
    other_accounts = [p.account for p in transaction.postings if p.account != cash_posting.account]
    
    for acc in other_accounts:
        # 投资活动相关账户
        if (acc.startswith('Assets:') and 
            ('Equipment' in acc or 'Property' in acc or 'Investment' in acc)):
            flow_type = 'investing'
            break
            
        # 筹资活动相关账户
        if (acc.startswith('Liabilities:Loans') or 
            acc.startswith('Equity:') or 
            'Capital' in acc):
            flow_type = 'financing'
            break
    
    return flow_type

def format_account_balances(accounts_dict) -> List[Dict]:
    """
    格式化账户余额为列表格式
    
    Args:
        accounts_dict: 账户余额字典
        
    Returns:
        格式化后的账户余额列表
    """
    result = []
    for account, currencies in accounts_dict.items():
        account_data = {"account": account, "balances": []}
        for currency, amount in currencies.items():
            account_data["balances"].append({
                "currency": currency,
                "amount": str(amount)
            })
        result.append(account_data)
    return result

def format_currency_amounts(currency_dict) -> List[Dict]:
    """
    格式化货币金额为列表格式
    
    Args:
        currency_dict: 货币金额字典
        
    Returns:
        格式化后的货币金额列表
    """
    result = []
    for currency, amount in currency_dict.items():
        result.append({
            "currency": currency,
            "amount": str(amount)
        })
    return result

def export_financial_statements(file_path: str = None, start_date: str = None, end_date: str = None) -> Dict:
    """
    导出完整的财务报表（资产负债表、利润表、现金流量表）
    
    Args:
        file_path: beancount文件路径
        start_date: 开始日期，格式为YYYY-MM-DD
        end_date: 结束日期，格式为YYYY-MM-DD
        
    Returns:
        包含三张财务报表的字典
    """
    try:
        # 生成三张报表
        balance_sheet = get_balance_sheet(file_path, end_date)
        income_statement = get_income_statement(file_path, start_date, end_date)
        cash_flow = get_cash_flow_statement(file_path, start_date, end_date)
        
        # 检查是否有错误
        if not balance_sheet.get("success", False):
            return {"success": False, "error": balance_sheet.get("error", "生成资产负债表失败")}
            
        if not income_statement.get("success", False):
            return {"success": False, "error": income_statement.get("error", "生成利润表失败")}
            
        if not cash_flow.get("success", False):
            return {"success": False, "error": cash_flow.get("error", "生成现金流量表失败")}
        
        # 合并结果
        return {
            "success": True,
            "period": {
                "start_date": start_date or "earliest",
                "end_date": end_date or "latest"
            },
            "balance_sheet": balance_sheet,
            "income_statement": income_statement,
            "cash_flow_statement": cash_flow
        }
    except Exception as e:
        return {"success": False, "error": str(e)}